// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once

#include "bfcallocator/common.h"
#include "bfcallocator/allocator.h"
#include <unordered_map>
#include <set>

namespace bfcallocator {

using AllocatorFactory = std::function<std::unique_ptr<IAllocator>(OrtDevice::DeviceId)>;

using AllocatorMap = std::unordered_map<int, AllocatorPtr>;
// TODO: update OrtMemoryInfo, use unordered_set instead
using MemoryInfoSet = std::set<OrtMemoryInfo>;

const int DEFAULT_CPU_ALLOCATOR_DEVICE_ID = 0;

struct AllocatorCreationInfo {
  AllocatorCreationInfo(AllocatorFactory device_alloc_factory,
                        OrtDevice::DeviceId device_id = 0,
                        bool use_arena = true,
                        OrtArenaCfg arena_cfg = {0, -1, -1, -1, -1})
      : device_alloc_factory(device_alloc_factory),
        device_id(device_id),
        use_arena(use_arena),
        arena_cfg(arena_cfg) {
  }

  AllocatorFactory device_alloc_factory;
  OrtDevice::DeviceId device_id;
  bool use_arena;
  OrtArenaCfg arena_cfg;
};

// Returns an allocator (an instance of IAllocator) based on the creation info provided.
// Returns nullptr if an invalid value of info.arena_cfg.arena_extend_strategy is supplied.
// Valid values can be found in bfcallocator_c_api.h.
AllocatorPtr CreateAllocator(const AllocatorCreationInfo& info);

#ifdef OPENGPU
class GPUAllocatorMgr
{
  public:
      static AllocatorPtr &GetInstance();

  private:
      GPUAllocatorMgr();
      ~GPUAllocatorMgr();
      GPUAllocatorMgr(const GPUAllocatorMgr &signal);
      const GPUAllocatorMgr &operator=(const GPUAllocatorMgr &signal);
      std::shared_ptr<IAllocator> allocator_ptr_;
};
#endif

class CPUAllocatorMgr
{
  public:
      static AllocatorPtr &GetInstance();

  private:
      CPUAllocatorMgr();
      ~CPUAllocatorMgr();
      CPUAllocatorMgr(const CPUAllocatorMgr &signal);
      const CPUAllocatorMgr &operator=(const CPUAllocatorMgr &signal);
      std::shared_ptr<IAllocator> allocator_ptr_;
};

}  // namespace bfcallocator
